#!/usr/bin/env python
# -*- coding:utf-8 -*-

import frida, sys
from madb.tools.printer import info,plus,warn
import time

class fhook:

    def __init__(self,device,attach) -> None:
        self.device = device
        self.attach = attach
    
    def start_poc(self,jscode):
        try:
            d = frida.get_device(self.device)
            session = d.attach(self.attach)
            pocScript  = session.create_script(jscode)
            info("hook程序开始...")
            pocScript.on('message', self.poc_message)
            pocScript.load()
            info(sys.stdin.read())
        except KeyboardInterrupt:
            plus("hook程序停止!")
        except Exception as e:
            print(e)
            warn("fhook错误!")
    def new_start_poc(self,jscode,func):
        try:
            d = frida.get_usb_device()
            session = d.attach(self.attach)
            pocScript  = session.create_script(jscode)
            info("hook程序开始...")
            pocScript.on('message', func)
            pocScript.load()
            info(sys.stdin.read())
        except KeyboardInterrupt:
            plus("hook程序停止!")
        except Exception as e:
            print(e)
            warn("fhook错误!")

    def call_rpc_func(self,jscode,func):
        d = frida.get_usb_device()
        session = d.attach(self.attach)
        pocScript  = session.create_script(jscode)
        info("hook程序开始...")
        pocScript.on('message', func)
        pocScript.load()
        return pocScript           
    def poc_message(self,message,data):
        if message['type'] == 'send':
            info(message['payload'])
        else:
            info(message)
